#管理训练和评估流程
"""
实验的初始化，包括模型、数据提供者和设备。
训练和评估迭代的实现。
模型保存和加载逻辑。
实验结果的记录和测试集评估。
"""

import tqdm
import os
import numpy as np
import sys

import wandb
from data import MetaLearningSystemDataLoader

from utils.storage import build_experiment_folder, save_statistics, save_to_json
import time
import torch


class ExperimentBuilder(object):
    def __init__(self, args, data, model, device):
        """
        Initializes an experiment builder using a named tuple (args), a data provider (data), a meta learning system
        (model) and a device (e.g. gpu/cpu/n)
        :param args: A namedtuple containing all experiment hyperparameters
        :param data: A data provider of instance MetaLearningSystemDataLoader
        :param model: A meta learning system instance
        :param device: Device/s to use for the experiment
        """
        self.args, self.device = args, device

        self.model = model
        self.saved_models_filepath, self.logs_filepath, self.samples_filepath = build_experiment_folder(
            experiment_name=self.args.experiment_name)

        self.data = data
        self.current_data_mode = "train"  # 当前激活的数据模式

        self.total_losses = dict()
        self.state = dict()
        self.state['best_val_acc'] = 0.
        self.state['best_val_iter'] = 0
        self.state['current_iter'] = 0
        self.state['current_iter'] = 0
        self.start_epoch = 0
        self.max_models_to_save = self.args.max_models_to_save
        self.create_summary_csv = False

        experiment_path = os.path.abspath(self.args.experiment_name)
        exp_name = experiment_path.split('/')[-1]
        log_base_dir = 'logs'
        os.makedirs(log_base_dir, exist_ok=True)

        log_dir = os.path.join(log_base_dir, exp_name)
        print(log_dir)

        if self.args.continue_from_epoch == 'from_scratch':
            self.create_summary_csv = True

        elif self.args.continue_from_epoch == 'latest':
            checkpoint = os.path.join(self.saved_models_filepath, "train_model_latest")
            print("attempting to find existing checkpoint", )
            if os.path.exists(checkpoint):
                self.state = \
                    self.model.load_model(model_save_dir=self.saved_models_filepath, model_name="train_model",
                                          model_idx='latest')
                self.start_epoch = int(self.state['current_iter'] / self.args.total_iter_per_epoch)

            else:
                self.args.continue_from_epoch = 'from_scratch'
                self.create_summary_csv = True
        elif int(self.args.continue_from_epoch) >= 0:
            self.state = \
                self.model.load_model(model_save_dir=self.saved_models_filepath, model_name="train_model",
                                      model_idx=self.args.continue_from_epoch)
            self.start_epoch = int(self.state['current_iter'] / self.args.total_iter_per_epoch)

        #self.data = data(args=args, current_iter=self.state['current_iter'])
        self.data.current_iter = self.state['current_iter']

        print("train_seed {}, val_seed: {}, at start time".format(self.data.dataset.seed["train"],
                                                                  self.data.dataset.seed["val"]))
        self.total_epochs_before_pause = self.args.total_epochs_before_pause
        self.state['best_epoch'] = int(self.state['best_val_iter'] / self.args.total_iter_per_epoch)
        self.epoch = int(self.state['current_iter'] / self.args.total_iter_per_epoch)
        self.augment_flag = True if 'omniglot' in self.args.train_dataset_name.lower() else False
        self.start_time = time.time()
        self.epochs_done_in_this_run = 0
        print(self.state['current_iter'], int(self.args.total_iter_per_epoch * self.args.total_epochs))

    def switch_data_mode(self, mode):
        """统一数据集切换接口"""
        self.current_data_mode = mode
        self.data.switch_set(mode)  # 调用数据集的切换方法

    def build_summary_dict(self, total_losses, phase, summary_losses=None):
        """
        Builds/Updates a summary dict directly from the metric dict of the current iteration.
        :param total_losses: Current dict with total losses (not aggregations) from experiment
        :param phase: Current training phase
        :param summary_losses: Current summarised (aggregated/summarised) losses stats means, stdv etc.
        :return: A new summary dict with the updated summary statistics information.
        """
        if summary_losses is None:
            summary_losses = dict()

        for key in total_losses:
            summary_losses["{}_{}_mean".format(phase, key)] = np.mean(total_losses[key])
            summary_losses["{}_{}_std".format(phase, key)] = np.std(total_losses[key])

        return summary_losses

    def build_loss_summary_string(self, summary_losses):
        """
        Builds a progress bar summary string given current summary losses dictionary
        :param summary_losses: Current summary statistics
        :return: A summary string ready to be shown to humans.
        """
        output_update = ""
        for key, value in zip(list(summary_losses.keys()), list(summary_losses.values())):
            if "loss" in key or "accuracy" in key:
                value = float(value)
                output_update += "{}: {:.4f}, ".format(key, value)

        return output_update

    def merge_two_dicts(self, first_dict, second_dict):
        """Given two dicts, merge them into a new dict as a shallow copy."""
        z = first_dict.copy()
        z.update(second_dict)
        return z

    def train_iteration(self, train_sample, sample_idx, epoch_idx, total_losses, current_iter, pbar_train):
        """
        Runs a training iteration, updates the progress bar and returns the total and current epoch train losses.
        :param train_sample: A sample from the data provider
        :param sample_idx: The index of the incoming sample, in relation to the current training run.
        :param epoch_idx: The epoch index.
        :param total_losses: The current total losses dictionary to be updated.
        :param current_iter: The current training iteration in relation to the whole experiment.
        :param pbar_train: The progress bar of the training.
        :return: Updates total_losses, train_losses, current_iter
        """
        x_support_set, x_target_set, y_support_set, y_target_set, seed = train_sample
        data_batch = (x_support_set, x_target_set, y_support_set, y_target_set)

        if sample_idx == 0:
            print("shape of data", x_support_set.shape, x_target_set.shape, y_support_set.shape,
                  y_target_set.shape)

        losses, _ = self.model.run_train_iter(data_batch=data_batch, epoch=epoch_idx)

        for key, value in zip(list(losses.keys()), list(losses.values())):
            if key not in total_losses:
                total_losses[key] = [float(value)]
            else:
                total_losses[key].append(float(value))

        train_losses = self.build_summary_dict(total_losses=total_losses, phase="train")
        train_output_update = self.build_loss_summary_string(losses)

        pbar_train.update(1)
        pbar_train.set_description("training phase {} -> {}".format(self.epoch, train_output_update))

        # Log training metrics to wandb
        if current_iter % 100 == 0:  # 每 100 次迭代记录一次
            wandb.log({"train_loss": losses["loss"], "train_accuracy": losses["accuracy"], "epoch": epoch_idx})

        current_iter += 1

        return train_losses, total_losses, current_iter

    def evaluation_iteration(self, val_sample, total_losses, pbar_val, phase):
        """
        Runs a validation iteration, updates the progress bar and returns the total and current epoch val losses.
        :param val_sample: A sample from the data provider
        :param total_losses: The current total losses dictionary to be updated.
        :param pbar_val: The progress bar of the val stage.
        :return: The updated val_losses, total_losses
        """
        x_support_set, x_target_set, y_support_set, y_target_set, seed = val_sample
        data_batch = (
            x_support_set, x_target_set, y_support_set, y_target_set)

        losses, _ = self.model.run_validation_iter(data_batch=data_batch)
        for key, value in zip(list(losses.keys()), list(losses.values())):
            if key not in total_losses:
                total_losses[key] = [float(value)]
            else:
                total_losses[key].append(float(value))

        val_losses = self.build_summary_dict(total_losses=total_losses, phase=phase)
        val_output_update = self.build_loss_summary_string(losses)

        pbar_val.update(1)
        pbar_val.set_description(
            "val_phase {} -> {}".format(self.epoch, val_output_update))

        # Log validation metrics to wandb
        wandb.log({"val_loss": losses["loss"], "val_accuracy": losses["accuracy"], "epoch": self.epoch})

        return val_losses, total_losses

    def test_evaluation_iteration(self, val_sample, model_idx, sample_idx, per_model_per_batch_preds, pbar_test):
        """
        Runs a validation iteration, updates the progress bar and returns the total and current epoch val losses.
        :param val_sample: A sample from the data provider
        :param total_losses: The current total losses dictionary to be updated.
        :param pbar_test: The progress bar of the val stage.
        :return: The updated val_losses, total_losses
        """
        x_support_set, x_target_set, y_support_set, y_target_set, seed = val_sample
        data_batch = (
            x_support_set, x_target_set, y_support_set, y_target_set)

        losses, per_task_preds = self.model.run_validation_iter(data_batch=data_batch)

        per_model_per_batch_preds[model_idx].extend(list(per_task_preds))

        test_output_update = self.build_loss_summary_string(losses)

        pbar_test.update(1)
        pbar_test.set_description(
            "test_phase {} -> {}".format(self.epoch, test_output_update))

        return per_model_per_batch_preds

    def save_models(self, model, epoch, state):
        """
        Saves two separate instances of the current model. One to be kept for history and reloading later and another
        one marked as "latest" to be used by the system for the next epoch training. Useful when the training/val
        process is interrupted or stopped. Leads to fault tolerant training and validation systems that can continue
        from where they left off before.
        :param model: Current meta learning model of any instance within the few_shot_learning_system.py
        :param epoch: Current epoch
        :param state: Current model and experiment state dict.
        """
        model.save_model(model_save_dir=os.path.join(self.saved_models_filepath, "train_model_{}".format(int(epoch))),
                         state=state)

        model.save_model(model_save_dir=os.path.join(self.saved_models_filepath, "train_model_latest"),
                         state=state)

        print("saved models to", self.saved_models_filepath)

    def pack_and_save_metrics(self, start_time, create_summary_csv, train_losses, val_losses, state):
        """
        Given current epochs start_time, train losses, val losses and whether to create a new stats csv file, pack stats
        and save into a statistics csv file. Return a new start time for the new epoch.
        :param start_time: The start time of the current epoch
        :param create_summary_csv: A boolean variable indicating whether to create a new statistics file or
        append results to existing one
        :param train_losses: A dictionary with the current train losses
        :param val_losses: A dictionary with the currrent val loss
        :return: The current time, to be used for the next epoch.
        """
        epoch_summary_losses = self.merge_two_dicts(first_dict=train_losses, second_dict=val_losses)

        if 'per_epoch_statistics' not in state:
            state['per_epoch_statistics'] = dict()

        for key, value in epoch_summary_losses.items():

            if key not in state['per_epoch_statistics']:
                state['per_epoch_statistics'][key] = [value]
            else:
                state['per_epoch_statistics'][key].append(value)

        epoch_summary_string = self.build_loss_summary_string(epoch_summary_losses)
        epoch_summary_losses["epoch"] = self.epoch
        epoch_summary_losses['epoch_run_time'] = time.time() - start_time

        if create_summary_csv:
            self.summary_statistics_filepath = save_statistics(self.logs_filepath, list(epoch_summary_losses.keys()),
                                                               create=True)
            self.create_summary_csv = False

        start_time = time.time()
        print("epoch {} -> {}".format(epoch_summary_losses["epoch"], epoch_summary_string))

        self.summary_statistics_filepath = save_statistics(self.logs_filepath,
                                                           list(epoch_summary_losses.values()))
        return start_time, state

    def evaluated_test_set_using_the_best_models(self, top_n_models):
        """
        使用验证集表现最好的前N个模型评估测试集
        """
        # 确保切换到测试模式
        self.data.switch_set("test")

        # 加载最佳模型
        per_epoch_stats = self.state['per_epoch_statistics']
        val_acc = np.array(per_epoch_stats['val_accuracy_mean'])
        top_indices = np.argsort(val_acc)[-top_n_models:][::-1]  # 取准确率最高的N个模型

        # 初始化存储结构
        all_preds = [[] for _ in range(top_n_models)]
        all_targets = [[] for _ in range(top_n_models)]

        # 对每个最佳模型进行测试评估
        for model_rank, model_idx in enumerate(top_indices):
            # 加载模型
            self.state = self.model.load_model(
                model_save_dir=self.saved_models_filepath,
                model_name="train_model",
                model_idx=model_idx + 1  # 假设模型按epoch+1命名
            )

            # 运行测试批次
            with tqdm.tqdm(total=int(self.args.num_evaluation_tasks / self.args.batch_size)) as pbar_test:
                for test_sample in self.data.get_test_batches(
                        total_batches=int(self.args.num_evaluation_tasks / self.args.batch_size),
                        augment_images=False
                ):
                    # 执行测试迭代
                    _, per_task_preds = self.model.run_validation_iter(test_sample)
                    all_preds[model_rank].extend(per_task_preds)
                    all_targets[model_rank].extend(test_sample[3].numpy())  # 假设目标在第四个位置
                    pbar_test.update(1)

        # 计算集成准确率
        ensemble_preds = np.mean(all_preds, axis=0)
        final_preds = np.argmax(ensemble_preds, axis=-1)
        targets = np.array(all_targets[0])  # 所有模型的目标相同

        accuracy = np.mean(final_preds == targets)
        accuracy_std = np.std(final_preds == targets)

        # 保存结果
        test_metrics = {
            'test_accuracy_mean': accuracy,
            'test_accuracy_std': accuracy_std
        }
        save_statistics(
            self.logs_filepath,
            list(test_metrics.keys()),
            create=True,
            filename="test_summary.csv"
        )
        save_statistics(
            self.logs_filepath,
            list(test_metrics.values()),
            create=False,
            filename="test_summary.csv"
        )
        print(f"测试结果: {test_metrics}")

    def run_experiment(self):
        """
        运行完整的训练实验，包含周期性的验证和最终的测试评估
        """
        # 初始化wandb
        wandb.init(project="MeTal", config=vars(self.args))

        # 确保初始模式为训练集
        self.data.switch_set("train")

        with tqdm.tqdm(initial=self.state['current_iter'],
                       total=int(self.args.total_iter_per_epoch * self.args.total_epochs)) as pbar_train:

            # 主训练循环
            while (self.state['current_iter'] < (self.args.total_epochs * self.args.total_iter_per_epoch)) and \
                    (not self.args.evaluate_on_test_set_only):

                # 训练阶段
                for train_sample in self.data.get_train_batches(
                        total_batches=int(self.args.total_iter_per_epoch * self.args.total_epochs) - self.state[
                            'current_iter'],
                        augment_images=self.augment_flag
                ):
                    # 执行单次训练迭代
                    train_losses, total_losses, self.state['current_iter'] = self.train_iteration(
                        train_sample=train_sample,
                        total_losses=self.total_losses,
                        epoch_idx=(self.state['current_iter'] // self.args.total_iter_per_epoch),
                        pbar_train=pbar_train,
                        current_iter=self.state['current_iter'],
                        sample_idx=self.state['current_iter']
                    )

                    # 检查是否完成一个epoch
                    if self.state['current_iter'] % self.args.total_iter_per_epoch == 0:
                        current_epoch = self.state['current_iter'] // self.args.total_iter_per_epoch

                        # 验证阶段
                        self.data.switch_set("val")  # 显式切换到验证模式
                        val_losses = {}
                        total_val_losses = {}
                        with tqdm.tqdm(total=int(self.args.num_evaluation_tasks / self.args.batch_size)) as pbar_val:
                            for val_sample in self.data.get_val_batches(
                                    total_batches=int(self.args.num_evaluation_tasks / self.args.batch_size),
                                    augment_images=False
                            ):
                                val_losses, total_val_losses = self.evaluation_iteration(
                                    val_sample=val_sample,
                                    total_losses=total_val_losses,
                                    pbar_val=pbar_val,
                                    phase='val'
                                )

                            # 更新最佳模型状态
                            if val_losses["val_accuracy_mean"] > self.state['best_val_acc']:
                                self.state.update({
                                    'best_val_acc': val_losses["val_accuracy_mean"],
                                    'best_val_iter': self.state['current_iter'],
                                    'best_epoch': current_epoch
                                })

                        # 保存模型和指标
                        self.epoch = current_epoch
                        self.state = self.merge_two_dicts(
                            self.merge_two_dicts(self.state, train_losses),
                            val_losses
                        )
                        self.save_models(self.model, self.epoch, self.state)

                        # 记录指标
                        self.start_time, self.state = self.pack_and_save_metrics(
                            self.start_time, self.create_summary_csv,
                            train_losses, val_losses, self.state
                        )

                        # 重置累计损失
                        self.total_losses = {}
                        self.epochs_done_in_this_run += 1

                        # 检查是否达到暂停条件
                        if self.epochs_done_in_this_run >= self.total_epochs_before_pause:
                            print(
                                f"训练种子 {self.data.dataset.seed['train']}, 验证种子 {self.data.dataset.seed['val']}，暂停前执行测试评估")
                            self.evaluated_test_set_using_the_best_models(top_n_models=5)
                            sys.exit()

                        # 切换回训练模式
                        self.data.switch_set("train")

        # 最终测试评估
        self.evaluated_test_set_using_the_best_models(top_n_models=5)

